#include "dns_base.h"
#define IS_PTR(os) (os >= 0xc000 && os <= 0xcfff)
#define GET_OFFSET(offset) (offset & 0x3fff)    //the 2 higher bits set to 0
#define SET_OFFSET(offset) (offset | 0xc000)



int  dns_set_name(dns_info_t* info, unsigned char* name,unsigned char* buff)
{
	int i = 0;
	unsigned short offset = 0;
	int len = 0;
	for (i = 0; i< info->header->q_count; i++)
	{
		if (strcmp(info->querys[i].qname, name) == 0)
		{
			offset = info->querys[i].qname - ((unsigned char*)info->header);
			offset =  SET_OFFSET(offset);
			break;
		}
	}

	for (i=0 ; i< info->header->ans_count+info->header->auth_count+info->header->add_count; i++)
	{
		if (strcmp(info->records[i].name, name) == 0)
		{
			offset = info->records[i].name - ((unsigned char*)info->header);
			offset = SET_OFFSET(offset);
			break;
		}
		if ( info->records[i].par->type==CNAME)
		{
			if (strcmp(info->records[i].rdata,name)==0)
			{
				offset = info->records[i].rdata - ((unsigned char*)info->header);
				offset = SET_OFFSET(offset);
				break;
			}
		}
	}

	if (offset == 0)
	{
		len = strlen(name) + 1;
		memcpy(buff, name, len);
	}
	else
	{
		offset = ntohs(offset);
		len = 2;
		memcpy(buff, &offset, 2);
	}
	return len;
}

int dns_get_name(unsigned char* buff, unsigned int start, unsigned int end, unsigned char** name)
{
	if (!buff || !name || start== end)
	{
		return 0;
	}
	unsigned char label_len;
	unsigned short offset = 0;
	unsigned short name_off = 0;
	unsigned char* itor = buff + start;
	int dlen = 0;
	int hasptr = 0, infinite = 20;
	unsigned char * labels[64];
	unsigned char   label_lens[64];
	unsigned char	label_index = 0;
	unsigned int    len = 0;
	while ((label_len=itor[0]) !=0  && (infinite--)) {
		offset = ntohs((unsigned short) * (unsigned short *)itor);
		if (IS_PTR(offset)) {
			itor = buff + GET_OFFSET(offset);
			if (hasptr == 0) {
				dlen = 2;
				dlen += len;
			}
			hasptr = 1;
			continue;
		}
		labels[label_index] = itor;
		label_lens[label_index] = itor[0]+1;  //3www|5baidu3com
		len += label_lens[label_index];
		label_index++;
		itor = itor + label_len + 1;
	}
	if (hasptr)
	{
		unsigned char * buff = (unsigned char*)malloc(len+1);//�� ��\0��
		(*name) = buff;
		unsigned char i = 0;
		for (; i < label_index; i++)
		{
			memcpy(buff, labels[i], label_lens[i]);
			buff += label_lens[i];
		}
		(*name)[len] = '\0';
	}
	else
	{
		(*name) = buff + start;
		len += 1; //增加1  结束符 \0
	}
		


	if (dlen == 0)
		dlen = len;          //root len is 1
	return dlen;
}


dns_info_t * create_dns_info(unsigned char * buff, unsigned short buff_len)
{
	if (!buff ||buff_len<=0 )
	{
		return NULL;
	}
	dns_info_t * info = malloc(sizeof(dns_info_t));
	if (!info)
		return NULL;
	memset(info, 0, sizeof(dns_info_t));
	info->buff = buff;
	info->buff_len = buff_len;
	info->start = 0;
	unsigned int len = dns_unmarshal_header(buff, 0, buff_len, &info->header);
	info->start += len;
	return info;
}

void free_dns_info(dns_info_t * info)
{
	if (!info)
	{
		return;
	}

	int i = 0;
	for (;i< info->header->q_count; i++)
	{
		if (!info->querys[i].qname)
		{
			continue;
		}
		if (info->querys[i].qname < info->buff || info->querys[i].qname > info->buff+info->buff_len)
		{
			free(info->querys[i].qname);
		}
	}
	int all_record = info->header->ans_count + info->header->auth_count + info->header->add_count;
	for (; i < all_record; i++)
	{
		if (!info->records[i].name)
		{
			if (info->records[i].name < info->buff || info->records[i].name > info->buff + info->buff_len)
			{
				free(info->records[i].name);
			}
		}

		if (!info->records[i].rdata)
		{
			if (info->records[i].rdata < info->buff || info->records[i].rdata > info->buff + info->buff_len)
			{
				free(info->records[i].rdata);
			}
		}
		
	}
	free(info);
}

//buff to struct
int dns_unmarshal_info(unsigned char* buff, unsigned int end, dns_info_t*info) {
	if (!buff || !info)
	{
		return 0;
	}

	int i = 0;
	if (info->start == 0)
	{
		info->start += dns_unmarshal_header(buff, 0, end, &info->header);
	}
	for (i = 0; i< info->header->q_count; i++)
	{
		info->start += dns_unmarshal_query(buff,info->start,end, &info->querys[i]);
	}

	int all_record = info->header->ans_count + info->header->auth_count + info->header->add_count;
	for (i = 0; i< all_record; i++)
	{
		info->start += dns_unmarshal_record(buff,info->start, end, &info->records[i]);
	}
}

int dns_unmarshal_header(unsigned char* buff, unsigned int start,unsigned int buff_len, dns_header_t** header) {

	if (NULL == buff || buff_len < DNS_HEADER_SIZE)
	{
		return 0;
	}
	*header = (dns_header_t*)buff;
	(*header)->id= ntohs((*header)->id);
	(*header)->q_count = ntohs((*header)->q_count);
	(*header)->ans_count = ntohs((*header)->ans_count);
	(*header)->auth_count = ntohs((*header)->auth_count);
	(*header)->add_count = ntohs((*header)->add_count);
	return DNS_HEADER_SIZE;
}

int dns_unmarshal_query(unsigned char* buff, unsigned int start, unsigned int buff_len, dns_query_t* query) {
	if (NULL == buff || buff_len < DNS_QUREY_MIN_SIZE)
	{
		return 0;
	}
	unsigned int len = dns_get_name(buff, start, buff_len, &query->qname);
	query->qname_len = strlen(query->qname) + 1;
	query->par = (dns_query_par_t*)(buff +start+ len);
	query->par->qclass = ntohs(query->par->qclass);
	query->par->qtype = ntohs(query->par->qtype);
	return len + sizeof(dns_query_par_t);
}

int dns_unmarshal_record(unsigned char* buff, unsigned int start, unsigned int buff_len, dns_record_t* record) {
	if (NULL == buff || buff_len < DNS_RECORD_MIN_SIZE)
	{
		return 0;
	}
	unsigned int len = dns_get_name(buff, start, buff_len,&record->name);
	record->name_len = strlen(record->name) + 1;
	record->par = (dns_record_par_t *)(buff + start + len);
	record->par->class = ntohs(record->par->class);
	record->par->type = ntohs(record->par->type);
	record->par->rdata_len = ntohs(record->par->rdata_len);
	record->par->ttl = ntohs(record->par->ttl);
	len += sizeof(dns_record_par_t);
	record->rdata = buff +start+ len  ;
	len += record->par->rdata_len;
	if (record->par->type == MX)
	{
		unsigned char * name = NULL;
		unsigned int rdata_len = 2;
		unsigned int start = (record->rdata - buff) + 2;
		dns_get_name(buff, start, record->par->rdata_len, &name);
		rdata_len += strlen(name) + 1;
		unsigned char *rdata = malloc(rdata_len);
		*(unsigned short*)rdata = *(unsigned short*)record->rdata;
		memcpy(rdata + 2, name, strlen(name) + 1);
		free(name);
		record->rdata = rdata;
		record->par->rdata_len = rdata_len;
	}
	else if (record->par->type !=A && record->par->type != AAAA)
	{
		unsigned char * name = NULL;
		dns_get_name(buff, record->rdata-buff, record->par->rdata_len, &name);
		record->rdata = name;
		record->par->rdata_len = strlen(name) + 1;
	}
	return len;
}




int dns_fill_header(dns_info_t *info) {
	if (NULL == info )
	{
		return 0;
	}
	info->header->id = htons(info->header->id);
	info->header->q_count = htons(info->header->q_count);
	info->header->ans_count = htons(info->header->ans_count);
	info->header->auth_count = htons(info->header->auth_count);
	info->header->add_count = htons(info->header->add_count);
	return 1;
}


int dns_fill_query(dns_info_t* info, unsigned char* qname, unsigned short qtype, unsigned short qclass){
	dns_query_t* query = &info->querys[info->header->q_count];
	
	query->qname = info->buff + info->start;
	query->qname_len = dns_set_name(info, qname, query->qname);
	query->par = query->qname + query->qname_len;
	query->par->qclass = htons(qclass);
	query->par->qtype = htons(qtype);
	info->header->q_count++;
	info->start += query->qname_len + sizeof(dns_query_par_t);
	return 1;
}



int dns_fill_record(dns_info_t* info, unsigned char* name, unsigned short type,unsigned short class,
					unsigned int ttl,unsigned short rdata_len, unsigned char * rdata) {
	int index = info->header->ans_count + info->header->auth_count + info->header->add_count;
	dns_record_t* record = &info->records[index];
	record->name = info->buff + info->start;
	record->name_len = dns_set_name(info, name, record->name);
	record->par = (dns_record_par_t *)(record->name + record->name_len);
	record->par->type = htons(type);
	record->par->ttl = htons(ttl);
	record->par->class = htons(class);
	
	record->rdata = record->name + record->name_len + sizeof(dns_record_par_t);
	if (type != A && type != MX && type != AAAA)
	{
		rdata_len  = dns_set_name(info, rdata, record->rdata);
	}
	else
	{
		memcpy(record->rdata, rdata, rdata_len);
	}
	record->par->rdata_len = htons(rdata_len);
	info->start += record->name_len + sizeof(dns_record_par_t)+ rdata_len;
	return 1;
}



int dns_fill_an(dns_info_t* info, unsigned char* name, unsigned short type, unsigned short class,
	unsigned int ttl, unsigned short rdata_len, unsigned char * rdata) {
	dns_fill_record(info,name,type,class,ttl,rdata_len,rdata);
	info->header->ans_count++;
	return 1;
}

int dns_fill_ns(dns_info_t* info, unsigned char* name, unsigned short type, unsigned short class,
	unsigned int ttl, unsigned short rdata_len, unsigned char * rdata) {
	dns_fill_record(info, name, type, class, ttl, rdata_len, rdata);
	info->header->auth_count++;
	return 1;
}
int dns_fill_au(dns_info_t* info, unsigned char* name, unsigned short type, unsigned short class,
	unsigned int ttl, unsigned short rdata_len, unsigned char * rdata) {

	dns_fill_record(info, name, type, class, ttl, rdata_len, rdata);
	info->header->add_count++;
	return 1;
}
